// topk
//block target, 求本 block 元素的 topk

#include <limits>

template<typename T>
__device__ void inner_sort(int end, T *src, T val, int stride=1){
  if(end == 0){
    src[0] = val;
  }else{
    T tmp;
    for(int k_id = end-1; k_id >= 0; --k_id){
      tmp = src[k_id*stride];
      // if(threadIdx.x == 0){
        // printf("val %f tmp %f k_id %d \n", val, tmp, k_id);
      // }
      if(val > tmp){
        if(k_id < (end-1)){
          src[(k_id+1)*stride] = tmp;
        }
        if(k_id == 0){
          src[0] = val;
        }
      }else{
        if(k_id < (end-1)){
          src[(k_id+1)*stride] = val;
        }
        break;
      }
    }
  }
}

template<typename T>
__global__ void topk_sm75(
  const int m, const int k, const T* __restrict__ src,
  T* __restrict__ dst, const int bt_sz, const T min_thres)
{
  const int tx = threadIdx.x;
  const int bx = blockIdx.x;
  const int szBlock = bt_sz*k;
  const int start = bx*szBlock;
  extern __shared__ T stmp[];

  // clear shared
  for(int id = tx; id < k*bt_sz; id+=bt_sz){
    stmp[id] = min_thres;
  }
  __syncthreads();
  for(int id = tx, k_id = 0; id < szBlock; id+=bt_sz, ++k_id){
    int idx = start+id;
    if(idx < m){
      inner_sort(k_id, stmp+tx, src[idx], bt_sz);
    }
  }
  __syncthreads();
  for(int stride = bt_sz>>1; stride > 0; stride >>=1){
    // for(int id = tx; id < k*bt_sz; id+=bt_sz){
      // printf(" (%d, %0.2f) ", id%bt_sz, stmp[id] < -99 ? -99 : stmp[id]);
    // }
    // __syncthreads();
    // if(tx == 0)
      // printf("\n///////////////////////////////////\n");
    // __syncthreads();    
    if(tx < stride){
      for(int k_id = 0; k_id < k; ++k_id){
        // if(tx == 0 && k_id == 0) printf("%0.2f \n", stmp[k_id*bt_sz+(tx+stride)]);
        inner_sort(k, stmp+tx, stmp[k_id*bt_sz+(tx+stride)], bt_sz);
      }
    }
    __syncthreads();
  }
  for(int k_id = tx; k_id < k; k_id+=bt_sz){
    dst[bx*k+k_id] = stmp[k_id*bt_sz];
  }
}

/**
 * TopK
 * k=20, m>80000, better cpu
*/
int topk(
  const int m, const int k, void *src, void *dst, const int bt_sz,
  cudaStream_t cuStream = nullptr)
{
  using SDT = float;
  using DDT = SDT;

  SDT *d_src = (SDT*)src;
  DDT *d_dst = (DDT*)dst;

  dim3 g_sz(DIV_UP(m, bt_sz*k)), b_sz(bt_sz);
  SDT min_thres = std::numeric_limits<SDT>::lowest();
  topk_sm75<<<g_sz, b_sz, bt_sz*k*sizeof(SDT), cuStream>>>(m, k, d_src, d_dst, bt_sz, min_thres);
  return 0;
}